
from pymongo import MongoClient, ASCENDING, DESCENDING, UpdateOne
from pandas import DataFrame
from pev2.util.DataBase import DB_CONN

"""
因子管理子系统
"""

class FactorModule:
    """
    获取某只股票的某个因子在一段时间内的值
    :param code:  股票代码
    :param factor:  因子名称
    :param startDate: 开始日期
    :param endDate: 结束日期
    :return: DataFrame(columns=['code',factor, 'date'])
    """
    def getSingleStockFactors(self, code, factor, startDate, endDate):
        factorConn = DB_CONN[factor]

        factorCursor = factorConn.find(
            {'code': code, 'date': {'$gte': startDate, '$lte': endDate}},
            sort=[('date', ASCENDING)]
        )

        factorDF = DataFrame(
            [{
                'date': x['date'],
                factor: x[factor],
                'code': x['code']
            }for x in factorCursor]
        )

        return factorDF

    """
    获取某个因子在某个交易日的所有股票的数据
    :param factor: 因子名称
    :param date: 日期
    :return: DataFrame(columns=['code',factor, 'date'])
    """
    def getSingleDateFactors(self, factor, date):
        factorConn = DB_CONN[factor]

        factorCursor = factorConn.find({'date': date})

        factorDF = DataFrame(
            [{
                'date': x['date'],
                factor: x[factor],
                'code': x['code']
            } for x in factorCursor]
        )

        return factorDF